import sys

import nmap

from network_scanner import ArgsParser, __version__, json_ret_wrapper, json_compress


class NetworkScanner:
    def __init__(self):
        self.__nmap = nmap.PortScanner()

    def _find_all_host(self, hosts, args):
        self.__nmap.scan(hosts=hosts, arguments=args)
        return self.__nmap.all_hosts()

    @json_compress
    @json_ret_wrapper
    def find_all_hosts(self, hosts: str, args: str = '-sP'):
        """
        find all HOSTs keeping alive in the subnet/host hosts.
        """
        print("uses the cmd: nmap -oX - %s" % args, file=sys.stderr)
        return self._find_all_host(hosts, args)

    def _find_alive_host_basic_info(self, alive_host: str, ports: str or None, args: str = '-sS -sV -O -v'):
        """
        find the basic info, such as mac, hostname, service version and os info etc., according the host alive specified.
        alive_host must be an ip.
        """

        if ports and 'U' in ports:
            args = args + ' -sU'

        result = self.__nmap.scan(hosts=alive_host, ports=ports, arguments=args)
        temp_rst = result['scan'][alive_host]

        ret = {'host': alive_host}

        if temp_rst['hostnames'][0]['name'] != '':
            ret['hostname'] = temp_rst['hostnames'][0]['name']
        if "mac" in temp_rst['addresses']:
            mac = temp_rst['addresses']['mac']
            ret['mac'] = mac
            if len(temp_rst['vendor']) != 0:
                ret['vendor'] = temp_rst['vendor'][mac]
        if temp_rst['status']['state'] == 'up':
            ret['status'] = temp_rst['status']['state']
        else:
            ret['status'] = temp_rst['status']['state']
            return ret

        if 'uptime' in temp_rst:
            ret['uptime'] = temp_rst['uptime']['seconds'] + 's'

        for proto in ('tcp', 'udp'):
            if proto in temp_rst:
                proto_dict = {}
                for port in temp_rst[proto]:
                    proto_dict['service'] = temp_rst[proto][port]['name']
                    if temp_rst[proto][port]['version'] != '':
                        proto_dict['service-version'] = temp_rst[proto][port]['version']
                    proto_dict['service-port'] = port
                    proto_dict['service-status'] = temp_rst[proto][port]['state']
                    if temp_rst[proto][port]['product'] != "":
                        proto_dict['service-product'] = temp_rst[proto][port]['product']
                ret[proto] = proto_dict
        if len(temp_rst['osmatch']) > 0:
            ret['os-name'] = temp_rst['osmatch'][0]['name']
            # cpe
            cpe = []
            for os in range(len(temp_rst['osmatch'][0]['osclass'])):
                cpe.append({
                    'type': temp_rst['osmatch'][0]['osclass'][os]['type'],
                    'cpe': temp_rst['osmatch'][0]['osclass'][os]['cpe'][0]
                })
            if len(cpe) > 0:
                ret['cpe'] = cpe
        return ret

    @json_compress
    @json_ret_wrapper
    def find_all_alive_hosts_basic_info(self, hosts, ports, args):
        alive_hosts = self._find_all_host(hosts, DEFAULT_FIND_HOSTS_ARGS)
        print("has found %s hosts alive in the subnet %s" % (len(alive_hosts), hosts), file=sys.stderr)
        print("uses the cmd: nmap -oX - %s" % args, file=sys.stderr)
        c = 1
        ret = []
        for alive_host in alive_hosts:
            ret.append(self._find_alive_host_basic_info(alive_host, ports, args))
            print("current progress is %s%%" % int(c * 100 / len(alive_hosts)), file=sys.stderr)
            c += 1
        return ret


DEFAULT_FIND_HOSTS_ARGS = "-sP"
DEFAULT_FIND_INFO_ARGS = "-Pn -n -sS -sV -O -v -T4"
DEFAULT_PORTS = None
DEFAULT_SUBNET = "127.0.0.1"


def main():
    parser = ArgsParser("")
    args, _ = parser.parse_arguments(sys.argv[1:])

    if args.version:
        print(__version__, file=sys.stdout)
        return

    net_scanner = NetworkScanner()

    if args.find_hosts:
        args_ = args.args or DEFAULT_FIND_HOSTS_ARGS
        print(net_scanner.find_all_hosts(args.find_hosts, args_), file=sys.stdout)

    if args.find_info:
        args_ = args.args or DEFAULT_FIND_INFO_ARGS
        ports = args.ports or DEFAULT_PORTS
        print(net_scanner.find_all_alive_hosts_basic_info(args.find_info, ports, args_), file=sys.stdout)


if __name__ == '__main__':
    main()
